{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to view and update past graph state\n",
    "\n",
    "Once you start [checkpointing](./persistence.ipynb) your graphs, you can easily\n",
    "**get** or **update** the state of the agent at any point in time. This permits\n",
    "a few things:\n",
    "\n",
    "1. You can surface a state during an interrupt to a user to let them accept an\n",
    "   action.\n",
    "2. You can **rewind** the graph to reproduce or avoid issues.\n",
    "3. You can **modify** the state to embed your agent into a larger system, or to\n",
    "   let the user better control its actions.\n",
    "\n",
    "The key methods used for this functionality are:\n",
    "\n",
    "- [getState](https://langgraph4j.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/CompiledGraph.html#getState-org.bsc.langgraph4j.RunnableConfig-):\n",
    "  fetch the values from the target config\n",
    "- [updateState](https://langgraph4j.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/CompiledGraph.html#updateState-org.bsc.langgraph4j.RunnableConfig-java.util.Map-java.lang.String-):\n",
    "  apply the given values to the target state\n",
    "\n",
    "**Note:** this requires passing in a checkpointer.\n",
    "\n",
    "This works for [StateGraph](https://langgraph4j.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/StateGraph.html)\n",
    "\n",
    "Below is an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "var userHomeDir = System.getProperty(\"user.home\");\n",
    "var localRespoUrl = \"file://\" + userHomeDir + \"/.m2/repository/\";\n",
    "var langchain4jVersion = \"1.0.1\";\n",
    "var langchain4jBeta = \"1.0.1-beta6\";\n",
    "var langgraph4jVersion = \"1.6-SNAPSHOT\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remove installed package from Jupiter cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash \n",
    "rm -rf \\{userHomeDir}/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/langgraph4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0mRepository \u001b[1m\u001b[32mlocal\u001b[0m url: \u001b[1m\u001b[32mfile:///Users/bsorrentino/.m2/repository/\u001b[0m added.\n",
      "\u001b[0mAdding dependency \u001b[0m\u001b[1m\u001b[32morg.slf4j:slf4j-jdk14:2.0.9\n",
      "\u001b[0mAdding dependency \u001b[0m\u001b[1m\u001b[32morg.bsc.langgraph4j:langgraph4j-core:1.6-SNAPSHOT\n",
      "\u001b[0mAdding dependency \u001b[0m\u001b[1m\u001b[32morg.bsc.langgraph4j:langgraph4j-langchain4j:1.6-SNAPSHOT\n",
      "\u001b[0mAdding dependency \u001b[0m\u001b[1m\u001b[32mdev.langchain4j:langchain4j:1.0.1\n",
      "\u001b[0mAdding dependency \u001b[0m\u001b[1m\u001b[32mdev.langchain4j:langchain4j-open-ai:1.0.1\n",
      "\u001b[0mAdding dependency \u001b[0m\u001b[1m\u001b[32mdev.langchain4j:langchain4j-ollama:1.0.1-beta6\n",
      "\u001b[0mSolving dependencies\n",
      "Resolved artifacts count: 17\n",
      "Add to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/slf4j/slf4j-jdk14/2.0.9/slf4j-jdk14-2.0.9.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/slf4j/slf4j-api/2.0.9/slf4j-api-2.0.9.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/langgraph4j/langgraph4j-core/1.6-SNAPSHOT/langgraph4j-core-1.6-SNAPSHOT.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/async/async-generator/3.2.2/async-generator-3.2.2.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/langgraph4j/langgraph4j-langchain4j/1.6-SNAPSHOT/langgraph4j-langchain4j-1.6-SNAPSHOT.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/dev/langchain4j/langchain4j/1.0.1/langchain4j-1.0.1.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/dev/langchain4j/langchain4j-core/1.0.1/langchain4j-core-1.0.1.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/jspecify/jspecify/1.0.0/jspecify-1.0.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/fasterxml/jackson/core/jackson-annotations/2.19.0/jackson-annotations-2.19.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/fasterxml/jackson/core/jackson-core/2.19.0/jackson-core-2.19.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/fasterxml/jackson/core/jackson-databind/2.19.0/jackson-databind-2.19.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/apache/opennlp/opennlp-tools/2.5.4/opennlp-tools-2.5.4.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/dev/langchain4j/langchain4j-open-ai/1.0.1/langchain4j-open-ai-1.0.1.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/dev/langchain4j/langchain4j-http-client/1.0.1/langchain4j-http-client-1.0.1.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/dev/langchain4j/langchain4j-http-client-jdk/1.0.1/langchain4j-http-client-jdk-1.0.1.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/knuddels/jtokkit/1.1.0/jtokkit-1.1.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/dev/langchain4j/langchain4j-ollama/1.0.1-beta6/langchain4j-ollama-1.0.1-beta6.jar\u001b[0m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "%dependency /add-repo local \\{localRespoUrl} release|never snapshot|always\n",
    "// %dependency /list-repos\n",
    "%dependency /add org.slf4j:slf4j-jdk14:2.0.9\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-core:\\{langgraph4jVersion}\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-langchain4j:\\{langgraph4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j:\\{langchain4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j-open-ai:\\{langchain4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j-ollama:\\{langchain4jBeta}\n",
    "\n",
    "%dependency /resolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Initialize logger**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "try( var file = new java.io.FileInputStream(\"./logging.properties\")) {\n",
    "    java.util.logging.LogManager.getLogManager().readConfiguration( file );\n",
    "}\n",
    "\n",
    "var log = org.slf4j.LoggerFactory.getLogger(\"time-travel\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the state\n",
    "\n",
    "State is an (immutable) data class, inheriting from prebuilt [MessagesState], shared with all nodes in our graph. A state is basically a wrapper of a `Map<String,Object>` that provides some enhancers:\n",
    "\n",
    "1. Schema (optional), that is a `Map<String,Channel>` where each [`Channel`] describe behaviour of the related property\n",
    "1. `value()` accessors that inspect Map an return an Optional of value contained and cast to the required type\n",
    "\n",
    "[`Channel`]: https://langgraph4j.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/state/Channel.html\n",
    "[MessagesState]: https://langgraph4j.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/prebuilt/MessagesState.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import org.bsc.langgraph4j.prebuilt.MessagesState;\n",
    "import org.bsc.langgraph4j.state.Channel;\n",
    "import dev.langchain4j.data.message.AiMessage;\n",
    "import dev.langchain4j.data.message.ChatMessage;\n",
    "\n",
    "public class State extends MessagesState<ChatMessage> {\n",
    "\n",
    "    public State(Map<String, Object> initData) {\n",
    "        super( initData  );\n",
    "    }\n",
    "\n",
    "\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the tools\n",
    "\n",
    "Using [langchain4j], We will first define the tools we want to use. For this simple example, we will\n",
    "use create a placeholder search engine. However, it is really easy to create\n",
    "your own tools - see documentation\n",
    "[here][tools] on how to do\n",
    "that.\n",
    "\n",
    "[langchain4j]: https://docs.langchain4j.dev\n",
    "[tools]: https://docs.langchain4j.dev/tutorials/tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dev.langchain4j.agent.tool.P;\n",
    "import dev.langchain4j.agent.tool.Tool;\n",
    "\n",
    "import java.util.Optional;\n",
    "\n",
    "import static java.lang.String.format;\n",
    "\n",
    "public class SearchTool {\n",
    "\n",
    "    @Tool(\"Use to surf the web, fetch current information, check the weather, and retrieve other information.\")\n",
    "    String execQuery(@P(\"The query to use in your search.\") String query) {\n",
    "\n",
    "        // This is a placeholder for the actual implementation\n",
    "        return \"Cold, with a low of 13 degrees\";\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the model\n",
    "\n",
    "Now we will load the\n",
    "[chat model].\n",
    "\n",
    "1. It should work with messages. We will represent all agent state in the form of messages, so it needs to be able to work well with them.\n",
    "2. It should work with [tool calling],meaning it can return function arguments in its response.\n",
    "\n",
    "Note:\n",
    "   >\n",
    "   > These model requirements are not general requirements for using LangGraph4j - they are just requirements for this one example.\n",
    "   >\n",
    "\n",
    "[chat model]: https://docs.langchain4j.dev/tutorials/chat-and-language-models\n",
    "[tool calling]: https://docs.langchain4j.dev/tutorials/tools   \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dev.langchain4j.model.ollama.OllamaChatModel;\n",
    "import dev.langchain4j.agent.tool.ToolSpecification;\n",
    "import dev.langchain4j.agent.tool.ToolSpecifications;\n",
    "\n",
    "var model = OllamaChatModel.builder()\n",
    "        .modelName( \"qwen2.5:7b\" )\n",
    "        .baseUrl(\"http://localhost:11434\")\n",
    "        .logResponses(true)\n",
    "        .maxRetries(2)\n",
    "        .temperature(0.0)\n",
    "        .build();\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test function calling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "execute: execQuery \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Optional[ToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import dev.langchain4j.agent.tool.ToolSpecification;\n",
    "import dev.langchain4j.agent.tool.ToolSpecifications;\n",
    "import dev.langchain4j.data.message.UserMessage;\n",
    "import dev.langchain4j.data.message.AiMessage;\n",
    "import dev.langchain4j.model.output.Response;\n",
    "import dev.langchain4j.model.openai.OpenAiChatModel;\n",
    "import dev.langchain4j.service.tool.DefaultToolExecutor;\n",
    "import org.bsc.langgraph4j.langchain4j.tool.LC4jToolService;\n",
    "import org.bsc.langgraph4j.langchain4j.serializer.std.LC4jStateSerializer;\n",
    "import dev.langchain4j.model.chat.request.ChatRequest;\n",
    "import dev.langchain4j.model.chat.request.ChatRequestParameters;\n",
    "\n",
    "var toolService = LC4jToolService.builder()\n",
    "                        .toolsFromObject( new SearchTool() )\n",
    "                        .build();\n",
    "\n",
    "var tools = toolService.toolSpecifications();\n",
    "\n",
    "UserMessage userMessage = UserMessage.from(\"What will the weather be like in London tomorrow?\");\n",
    "\n",
    "var params = ChatRequestParameters.builder()\n",
    "                .toolSpecifications( tools )\n",
    "                .build();\n",
    "var request = ChatRequest.builder()\n",
    "                .parameters( params )\n",
    "                .messages( userMessage )\n",
    "                .build();\n",
    "\n",
    "var response = model.chat( request );\n",
    "\n",
    "var result = toolService.execute( response.aiMessage().toolExecutionRequests() );\n",
    "\n",
    "result;\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the graph\n",
    "\n",
    "We can now put it all together. We will run it first without a checkpointer:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import static org.bsc.langgraph4j.StateGraph.START;\n",
    "import static org.bsc.langgraph4j.StateGraph.END;\n",
    "import static org.bsc.langgraph4j.action.AsyncEdgeAction.edge_async;\n",
    "import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async;\n",
    "import org.bsc.langgraph4j.StateGraph;\n",
    "import org.bsc.langgraph4j.action.EdgeAction;\n",
    "import org.bsc.langgraph4j.action.NodeAction;\n",
    "import dev.langchain4j.data.message.AiMessage;\n",
    "import dev.langchain4j.data.message.ChatMessage;\n",
    "import dev.langchain4j.service.tool.DefaultToolExecutor;\n",
    "import org.bsc.langgraph4j.checkpoint.MemorySaver; \n",
    "import org.bsc.langgraph4j.CompileConfig; \n",
    "import java.util.stream.Collectors;\n",
    "\n",
    "// Route Message \n",
    "EdgeAction<State> routeMessage = state -> {\n",
    "  \n",
    "  var lastMessage = state.lastMessage();\n",
    "  \n",
    "  if ( !lastMessage.isPresent()) return \"exit\";\n",
    "\n",
    "  if( lastMessage.get() instanceof AiMessage message  ) {\n",
    "\n",
    "    // If tools should be called\n",
    "    if ( message.hasToolExecutionRequests() ) return \"next\";\n",
    "    \n",
    "  }\n",
    "  \n",
    "  // If no tools are called, we can finish (respond to the user)\n",
    "  return \"exit\";\n",
    "};\n",
    "\n",
    "// Call Model\n",
    "NodeAction<State> callModel = state -> {\n",
    "  var tools = ToolSpecifications.toolSpecificationsFrom( SearchTool.class );\n",
    "\n",
    "  var params = ChatRequestParameters.builder()\n",
    "                .toolSpecifications( tools )\n",
    "                .build();\n",
    "  var request = ChatRequest.builder()\n",
    "                .parameters( params )\n",
    "                .messages( state.messages() )\n",
    "                .build();\n",
    "\n",
    "  var response = model.chat( request );\n",
    "  \n",
    "\n",
    "  return Map.of( \"messages\", response.aiMessage() );\n",
    "};\n",
    "\n",
    "final var toolService = LC4jToolService.builder()\n",
    "                      .toolsFromObject(new SearchTool())\n",
    "                      .build();\n",
    "// Invoke Tool \n",
    "NodeAction<State> invokeTool = state -> {\n",
    "  var lastMessage = (AiMessage)state.lastMessage()\n",
    "                          .orElseThrow( () -> ( new IllegalStateException( \"last message not found!\")) );\n",
    "  \n",
    "  var result = toolService.execute( lastMessage.toolExecutionRequests() )\n",
    "                    .orElseThrow( () -> ( new IllegalStateException( \"tool execution failed!\")));\n",
    "\n",
    "  return Map.of( \"messages\", result );\n",
    "};\n",
    "\n",
    "var stateSerializer = new LC4jStateSerializer<>(State::new);\n",
    "\n",
    "// Define Graph\n",
    "var workflow = new StateGraph<State>(State.SCHEMA, stateSerializer)\n",
    "  .addNode(\"agent\", node_async(callModel) )\n",
    "  .addNode(\"tools\", node_async(invokeTool) )\n",
    "  .addEdge(START, \"agent\")\n",
    "  .addConditionalEdges(\"agent\", edge_async(routeMessage), Map.of( \"next\", \"tools\", \"exit\", END ))\n",
    "  .addEdge(\"tools\", \"agent\");\n",
    "\n",
    "// Here we only save in-memory\n",
    "var memory = new MemorySaver();\n",
    "\n",
    "var compileConfig = CompileConfig.builder()\n",
    "                    .checkpointSaver(memory)\n",
    "                    .releaseThread(false) // DON'T release thread after completion\n",
    "                    .build();\n",
    "\n",
    "var graph = workflow.compile(compileConfig);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interacting with the Agent\n",
    "\n",
    "We can now interact with the agent. Between interactions you can get and update state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "START \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__START__\n",
      "{\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"Hi I'm Bartolo.\" }] }\n",
      "\t]\n",
      "\t}\n",
      "agent\n",
      "{\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"Hi I'm Bartolo.\" }] }\n",
      "\tAiMessage { text = \"Hello Bartolo! Nice to meet you. How can I assist you today?\" toolExecutionRequests = [] }\n",
      "\t]\n",
      "\t}\n",
      "__END__\n",
      "{\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"Hi I'm Bartolo.\" }] }\n",
      "\tAiMessage { text = \"Hello Bartolo! Nice to meet you. How can I assist you today?\" toolExecutionRequests = [] }\n",
      "\t]\n",
      "\t}\n"
     ]
    }
   ],
   "source": [
    "import org.bsc.langgraph4j.RunnableConfig;\n",
    "\n",
    "var runnableConfig =  RunnableConfig.builder()\n",
    "                .threadId(\"conversation-num-1\" )\n",
    "                .build();\n",
    "\n",
    "Map<String,Object> inputs = Map.of( \"messages\", UserMessage.from(\"Hi I'm Bartolo.\") );\n",
    "\n",
    "var result = graph.stream( inputs, runnableConfig );\n",
    "\n",
    "for( var r : result ) {\n",
    "  System.out.println( r.node() );\n",
    "  System.out.println( r.state() );\n",
    "  \n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here you can see the \"`agent`\" node ran, and then our edge returned `__END__` so the graph stopped execution there.\n",
    "\n",
    "Let's check the current graph state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "StateSnapshot{node=agent, state={\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"Hi I'm Bartolo.\" }] }\n",
      "\tAiMessage { text = \"Hello Bartolo! Nice to meet you. How can I assist you today?\" toolExecutionRequests = [] }\n",
      "\t]\n",
      "\t}, config=RunnableConfig{ threadId=conversation-num-1, checkPointId=18dbbc9f-a203-45d4-9e55-8f0f6622c2ac, nextNode=__END__, streamMode=VALUES }}\n"
     ]
    }
   ],
   "source": [
    "import org.bsc.langgraph4j.checkpoint.Checkpoint;\n",
    "\n",
    "var checkpoint = graph.getState(runnableConfig);\n",
    "\n",
    "System.out.println(checkpoint);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The current state is the two messages we've seen above, 1. the Human Message we sent in, 2. the AIMessage we got back from the model.\n",
    "\n",
    "The next value is `__END__`  since the graph has terminated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "__END__"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "checkpoint.getNext()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's get it to execute a tool\n",
    "\n",
    "When we call the graph again, it will create a checkpoint after each internal execution step. Let's get it to run a tool, then look at the checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "START \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "execute: execQuery \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AiMessage { text = \"The current weather in San Francisco is quite cold, with temperatures at around 13 degrees. Please make sure to dress warmly!\" toolExecutionRequests = [] }\n"
     ]
    }
   ],
   "source": [
    "\n",
    "Map<String,Object> inputs = Map.of( \"messages\", UserMessage.from(\"What's the weather like in SF currently?\") );\n",
    "\n",
    "var state = graph.invoke( inputs, runnableConfig ).orElseThrow( () ->(new IllegalStateException()) ) ;\n",
    "\n",
    "System.out.println( state.lastMessage().orElse(null) );\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pause before tools\n",
    "\n",
    "If you notice below, we now will add interruptBefore=[\"action\"] - this means that before any actions are taken we pause. This is a great moment to allow the user to correct and update the state! This is very useful when you want to have a human-in-the-loop to validate (and potentially change) the action to take.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "START \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionRequest id is null! \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__START__\n",
      "{\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\t]\n",
      "\t}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ToolExecutionRequest id is null! \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "agent\n",
      "{\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\tAiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"query\" : \"current weather in San Francisco\"\n",
      "}\" }] }\n",
      "\t]\n",
      "\t}\n"
     ]
    }
   ],
   "source": [
    "var memory = new MemorySaver();\n",
    "\n",
    "var compileConfig = CompileConfig.builder()\n",
    "                    .checkpointSaver(memory)\n",
    "                    .releaseThread(false) // DON'T release thread after completion\n",
    "                    .interruptBefore( \"tools\")\n",
    "                    .build();\n",
    "\n",
    "var graphWithInterrupt = workflow.compile(compileConfig);\n",
    "\n",
    "var runnableConfig =  RunnableConfig.builder()\n",
    "                .threadId(\"conversation-2\" )\n",
    "                .build();\n",
    "\n",
    "Map<String,Object> inputs = Map.of( \"messages\", UserMessage.from(\"What's the weather like in SF currently?\") );\n",
    "\n",
    "var result = graphWithInterrupt.stream( inputs, runnableConfig );\n",
    "\n",
    "for( var r : result ) {\n",
    "  System.out.println( r.node() );\n",
    "  System.out.println( r.state() );\n",
    "  \n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get State\n",
    "\n",
    "You can fetch the latest graph checkpoint using `getState(config)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tools"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var snapshot = graphWithInterrupt.getState(runnableConfig);\n",
    "snapshot.getNext();\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resume\n",
    "\n",
    "You can resume by running the graph with a null input. The checkpoint is loaded, and with no new inputs, it will execute as if no interrupt had occurred."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "RESUME REQUEST \n",
      "RESUME FROM agent \n",
      "ToolExecutionRequest id is null! \n",
      "execute: execQuery \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "RESULT:\n",
      "tools\n",
      "{\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\tAiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"query\" : \"current weather in San Francisco\"\n",
      "}\" }] }\n",
      "\tToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }\n",
      "\t]\n",
      "\t} \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "RESULT:\n",
      "agent\n",
      "{\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\tAiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"query\" : \"current weather in San Francisco\"\n",
      "}\" }] }\n",
      "\tToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }\n",
      "\tAiMessage { text = \"The current weather in San Francisco is cold, with temperatures currently at 13 degrees. Please make sure to dress warmly!\" toolExecutionRequests = [] }\n",
      "\t]\n",
      "\t} \n",
      "RESULT:\n",
      "__END__\n",
      "{\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\tAiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"query\" : \"current weather in San Francisco\"\n",
      "}\" }] }\n",
      "\tToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }\n",
      "\tAiMessage { text = \"The current weather in San Francisco is cold, with temperatures currently at 13 degrees. Please make sure to dress warmly!\" toolExecutionRequests = [] }\n",
      "\t]\n",
      "\t} \n"
     ]
    }
   ],
   "source": [
    "import org.bsc.langgraph4j.GraphInput;\n",
    "\n",
    "var result = graphWithInterrupt.stream( GraphInput.resume(), snapshot.getConfig() );\n",
    "\n",
    "for( var r : result ) {\n",
    "  log.trace( \"RESULT:\\n{}\\n{}\", r.node(), r.state() );\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check full history\n",
    "\n",
    "Let's browse the history of this thread, from newest to oldest."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "---\n",
      "StateSnapshot{node=agent, state={\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\tAiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"query\" : \"current weather in San Francisco\"\n",
      "}\" }] }\n",
      "\tToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }\n",
      "\tAiMessage { text = \"The current weather in San Francisco is cold, with temperatures currently at 13 degrees. Please make sure to dress warmly!\" toolExecutionRequests = [] }\n",
      "\t]\n",
      "\t}, config=RunnableConfig{ threadId=conversation-2, checkPointId=1ee4643d-e2ba-4143-8685-9d284e24cb0e, nextNode=__END__, streamMode=VALUES }}\n",
      "--- \n",
      "\n",
      "---\n",
      "StateSnapshot{node=tools, state={\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\tAiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"query\" : \"current weather in San Francisco\"\n",
      "}\" }] }\n",
      "\tToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }\n",
      "\t]\n",
      "\t}, config=RunnableConfig{ threadId=conversation-2, checkPointId=d5fefcc9-a64c-4902-a653-851a2b4f77f0, nextNode=agent, streamMode=VALUES }}\n",
      "--- \n",
      "\n",
      "---\n",
      "StateSnapshot{node=agent, state={\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\tAiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"query\" : \"current weather in San Francisco\"\n",
      "}\" }] }\n",
      "\t]\n",
      "\t}, config=RunnableConfig{ threadId=conversation-2, checkPointId=b49955fc-5c59-47ec-9087-e4f3d62c91fd, nextNode=tools, streamMode=VALUES }}\n",
      "--- \n",
      "\n",
      "---\n",
      "StateSnapshot{node=__START__, state={\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\t]\n",
      "\t}, config=RunnableConfig{ threadId=conversation-2, checkPointId=6d6d7e39-76af-4c16-804d-721f71615a9f, nextNode=agent, streamMode=VALUES }}\n",
      "--- \n"
     ]
    }
   ],
   "source": [
    "RunnableConfig toReplay = null;\n",
    "var states = graphWithInterrupt.getStateHistory(runnableConfig);\n",
    "for( var state: states ) {\n",
    "  \n",
    "  log.trace( \"\\n---\\n{}\\n---\",state);\n",
    "\n",
    "  if (state.state().messages().size() == 3) {\n",
    "     toReplay = state.getConfig();\n",
    "  }\n",
    "}\n",
    "if (toReplay==null) {\n",
    "  throw new IllegalStateException(\"No state to replay\");\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Replay a past state\n",
    "\n",
    "To replay from this place we just need to pass its config back to the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "RESUME REQUEST \n",
      "RESUME FROM tools \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "ToolExecutionRequest id is null! \n",
      "ToolExecutionResultMessage id is null! \n",
      "RESULT:\n",
      "agent\n",
      "{\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\tAiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"query\" : \"current weather in San Francisco\"\n",
      "}\" }] }\n",
      "\tToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }\n",
      "\tAiMessage { text = \"The current weather in San Francisco is cold, with temperatures currently at 13 degrees. Please make sure to dress warmly!\" toolExecutionRequests = [] }\n",
      "\t]\n",
      "\t}\n",
      "--- \n",
      "RESULT:\n",
      "__END__\n",
      "{\n",
      "\tmessages=[\n",
      "\tUserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }\n",
      "\tAiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = null, name = \"execQuery\", arguments = \"{\n",
      "  \"query\" : \"current weather in San Francisco\"\n",
      "}\" }] }\n",
      "\tToolExecutionResultMessage { id = null toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }\n",
      "\tAiMessage { text = \"The current weather in San Francisco is cold, with temperatures currently at 13 degrees. Please make sure to dress warmly!\" toolExecutionRequests = [] }\n",
      "\t]\n",
      "\t}\n",
      "--- \n"
     ]
    }
   ],
   "source": [
    "var results = graphWithInterrupt.stream( GraphInput.resume(), toReplay ); \n",
    "\n",
    "for( var r : results ) {\n",
    "  log.trace( \"RESULT:\\n{}\\n{}\\n---\", r.node(), r.state() );\n",
    "}\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java (rjk 2.2.0)",
   "language": "java",
   "name": "rapaio-jupyter-kernel"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "java",
   "nbconvert_exporter": "script",
   "pygments_lexer": "java",
   "version": "22.0.2+9-70"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
